install.packages("survival", dependencies=TRUE)
install.packages("mstate", dependencies=TRUE)

rm(list=ls())
library(foreign)
library(survival)
library(mstate)

#Make Sure Replication Data is in the Same Folder as This File
dat.1 <- read.dta("use_rep.dta")

#Drop South Africa. See code in Appendix for model estimates including South Africa.
dat.1 <- subset(dat.1, dat.1$ccode != 560)

#####################################
##### Prepare Transition Matrix #####
#####################################

tmat <- transMat(list(c(2), #transitions from no program
			c(1, 3), 		#transitions from program
			c()),			#transition from weapon			
			names = c("no_program", "program", "weapon"))

##### Convert Data Type
attr(dat.1, "trans") <- tmat
class(dat.1) <- c("msdata", "data.frame")


############################
######## MODELS ############
############################

#Replicate Table 2 in Main Text
mod.1 <- coxph(Surv(t, status) ~ 
		nuk7set1_1 + npt_eff_1+   lmid_past5yrs_1 +  neopatr_1+
		nuk7set1_2 + npt_eff_2+   lmid_past5yrs_2+ neopatr_2+
		nuk7set1_3 + npt_eff_3 +  lmid_past5yrs_3+ neopatr_3
				+ strata(trans), data = dat.1, method = "breslow")
summary(mod.1)



#############################
######## Figures ############
#############################

##### Replicate Figure 2 #####
#Estimate Nonparametric Multi-State Model
np.1 <- coxph(Surv(t, status) ~ strata(trans), data = dat.1, method = "breslow")

#Estimate Cumulative Hazards
msf1 <- msfit(object = np.1, vartype = "aalen", trans = tmat)

#Simulate Transition Probabilities
set.seed(1009)
np.samp <- mssample(Haz = msf1$Haz, trans = tmat, tvec = 0:30, history=list(state=2, time=0), M = 2500, clock = "reset", output = "state")

#Plot Simulated Transition Probabilities
plot(np.samp$time, np.samp$pstate1, type="n" ,xlim=c(0,30), ylim=c(0,1), xlab="Years Since Program Initiated", ylab="Probability", main="")
lines(np.samp$pstate1, lty=1)
lines(np.samp$pstate2, lty=2)
lines(np.samp$pstate3, lty=3)
legend(7,.9, c("No Nuclear Program", "Nuclear Program","Nuclear Weapon"), lty = c(1,2,3), bty="n", cex=.8)



##### Replicate Figure 3 #####

#Define Covariate Values
neopatr_1 <- as.vector(rep(0, 3))
neopatr_1[1] <- mean(dat.1$neopatr, na.rm=TRUE)
neopatr_2 <- as.vector(rep(0, 3))
neopatr_2[2] <- mean(dat.1$neopatr, na.rm=TRUE)
neopatr_3 <- as.vector(rep(0, 3))
neopatr_3[3] <- mean(dat.1$neopatr, na.rm=TRUE)

lmid_past5yrs_1 <- as.vector(rep(0, 3))
lmid_past5yrs_1[1] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)
lmid_past5yrs_2 <- as.vector(rep(0, 3))
lmid_past5yrs_2[2] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)
lmid_past5yrs_3 <- as.vector(rep(0, 3))
lmid_past5yrs_3[3] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)

nuk7set1_1 <- as.vector(rep(0, 3))
nuk7set1_1[1] <- mean(dat.1$nuk7set1, na.rm=TRUE)
nuk7set1_2 <- as.vector(rep(0, 3))
nuk7set1_2[2] <- mean(dat.1$nuk7set1, na.rm=TRUE)
nuk7set1_3 <- as.vector(rep(0, 3))
nuk7set1_3[3] <- mean(dat.1$nuk7set1, na.rm=TRUE)

#Set NPT Membership to High
npt_eff_1 <- as.vector(rep(0, 3))
npt_eff_1[1] <- quantile(dat.1$npt_eff, na.rm=TRUE, .75)
npt_eff_2 <- as.vector(rep(0, 3))
npt_eff_2[2] <- quantile(dat.1$npt_eff, na.rm=TRUE, .75)
npt_eff_3 <- as.vector(rep(0, 3))
npt_eff_3[3] <- quantile(dat.1$npt_eff, na.rm=TRUE, .75)

sit.hinpt <- data.frame(cbind(
							nuk7set1_1, lmid_past5yrs_1, npt_eff_1, neopatr_1,
							nuk7set1_2, lmid_past5yrs_2, npt_eff_2, neopatr_2,
							nuk7set1_3, lmid_past5yrs_3, npt_eff_3, neopatr_3))
							
sit.hinpt$trans <- 1:3							
attr(sit.hinpt, "trans") <- tmat
class(sit.hinpt) <- c("msdata", "data.frame")						
sit.hinpt$strata <- sit.hinpt$trans

#Estimate Cumulative Hazards When NPT Membership is High
msf.hinpt <- msfit(mod.1, sit.hinpt, trans=tmat, variance=FALSE)

#Simulate Transition Probabilities When NPPT Membership is High
set.seed(1009)
ms.hinpt <- mssample(Haz = msf.hinpt$Haz, trans = tmat, tvec = 0:40, history=list(state=1, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)

#Set NPT Membership to Low
npt_eff_1 <- as.vector(rep(0, 3))
npt_eff_1[1] <- quantile(dat.1$npt_eff, na.rm=TRUE, .25)
npt_eff_2 <- as.vector(rep(0, 3))
npt_eff_2[2] <- quantile(dat.1$npt_eff, na.rm=TRUE, .25)
npt_eff_3 <- as.vector(rep(0, 3))
npt_eff_3[3] <- quantile(dat.1$npt_eff, na.rm=TRUE, .25)

sit.lonpt <- data.frame(cbind(
							nuk7set1_1, lmid_past5yrs_1, npt_eff_1, neopatr_1,
							nuk7set1_2, lmid_past5yrs_2, npt_eff_2, neopatr_2,
							nuk7set1_3, lmid_past5yrs_3, npt_eff_3, neopatr_3))
							
sit.lonpt$trans <- 1:3							
attr(sit.lonpt, "trans") <- tmat
class(sit.lonpt) <- c("msdata", "data.frame")						
sit.lonpt$strata <- sit.lonpt$trans

#Estimate Cumulative Hazards When NPT Membership is Low		
msf.lonpt <- msfit(mod.1, sit.lonpt, trans=tmat, variance=FALSE)

#Simulate Transition Probabilities When NPPT Membership is Low
set.seed(1009)
ms.lonpt <- mssample(Haz = msf.lonpt$Haz, trans = tmat, tvec = 0:40, history=list(state=1, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)

#Plot Simulated Transition Probabilities For Low and High NPT Membership
plot(ms.lonpt$time, ms.lonpt$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,.2), xlab="Years", ylab="Probability of Nuclear Program ", main="")
lines(ms.lonpt$time, ms.lonpt$pstate2, lty=1)
lines(ms.hinpt$time, ms.hinpt$pstate2, lty=2)
legend(3,.18, c("Low NPT Membership","High NPT Membership"), lty = c(1,2), bty="n")




##### Replicate Figure 4 #####

#Simulate Transition Probabilities When NPPT Membership is High, vary starting phase to nuclear program
set.seed(1009)
ms.hinpt.2 <- mssample(Haz = msf.hinpt$Haz, trans = tmat, tvec = 0:40, history=list(state=2, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)

##Simulate Transition Probabilities When NPPT Membership is Low, vary starting phase to nuclear program
set.seed(1009)
ms.lonpt.2 <- mssample(Haz = msf.lonpt$Haz, trans = tmat, tvec = 0:40, history=list(state=2, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)

#Figure 4a
plot(ms.lonpt.2$time, ms.lonpt.2$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,1), xlab="Years Since Program Initiated", ylab="Probability of Nuclear Program", main="")
lines(ms.lonpt.2$time, ms.lonpt.2$pstate2, lty=1)
lines(ms.hinpt.2$time, ms.hinpt.2$pstate2, lty=2)
legend(3,.4, c("Low NPT Membership","High NPT Membership"), lty = c(1,2), bty="n")

#Figure 4b
plot(ms.lonpt.2$time, ms.lonpt.2$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,.3), xlab="Years Since Program Initiated", ylab="Probability of No Nuclear Program or Weapon", main="")
lines(ms.lonpt.2$time, ms.lonpt.2$pstate1, lty=1)
lines(ms.hinpt.2$time, ms.hinpt.2$pstate1, lty=2)
legend(3,.2, c("Low NPT Membership","High NPT Membership"), lty = c(1,2), bty="n")

plot(ms.lonpt.2$time, ms.lonpt.2$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,.008), xlab="Years Since Program Initiated", ylab="Probability of Nuclear Weapon", main="")
lines(ms.lonpt.2$time, ms.lonpt.2$pstate3, lty=1)
lines(ms.hinpt.2$time, ms.hinpt.2$pstate3, lty=2)
legend(-.8,.004, c("Low NPT Membership","High NPT Membership"), lty = c(1,2), bty="n")





##### Replicate Figure 5 #####

#Define Covariate Values
neopatr_1 <- as.vector(rep(0, 3))
neopatr_1[1] <- mean(dat.1$neopatr, na.rm=TRUE)
neopatr_2 <- as.vector(rep(0, 3))
neopatr_2[2] <- mean(dat.1$neopatr, na.rm=TRUE)
neopatr_3 <- as.vector(rep(0, 3))
neopatr_3[3] <- mean(dat.1$neopatr, na.rm=TRUE)

npt_eff_1 <- as.vector(rep(0, 3))
npt_eff_1[1] <- mean(dat.1$npt_eff, na.rm=TRUE)
npt_eff_2 <- as.vector(rep(0, 3))
npt_eff_2[2] <- mean(dat.1$npt_eff, na.rm=TRUE)
npt_eff_3 <- as.vector(rep(0, 3))
npt_eff_3[3] <- mean(dat.1$npt_eff, na.rm=TRUE)

lmid_past5yrs_1 <- as.vector(rep(0, 3))
lmid_past5yrs_1[1] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)
lmid_past5yrs_2 <- as.vector(rep(0, 3))
lmid_past5yrs_2[2] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)
lmid_past5yrs_3 <- as.vector(rep(0, 3))
lmid_past5yrs_3[3] <- mean(dat.1$lmid_past5yrs, na.rm=TRUE)

#Set nuclear latency to high
nuk7set1_1 <- as.vector(rep(0, 3))
nuk7set1_1[1] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .75)
nuk7set1_2 <- as.vector(rep(0, 3))
nuk7set1_2[2] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .75)
nuk7set1_3 <- as.vector(rep(0, 3))
nuk7set1_3[3] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .75)

sit.hicap <- data.frame(cbind(
							nuk7set1_1, lmid_past5yrs_1, npt_eff_1, neopatr_1,
							nuk7set1_2, lmid_past5yrs_2, npt_eff_2, neopatr_2,
							nuk7set1_3, lmid_past5yrs_3, npt_eff_3, neopatr_3))
				
sit.hicap$trans <- 1:3							
attr(sit.hicap, "trans") <- tmat
class(sit.hicap) <- c("msdata", "data.frame")						
sit.hicap$strata <- sit.hicap$trans

#Estimate Cumulative Hazards When latency is High
msf.hicap <- msfit(mod.1, sit.hicap, trans=tmat, variance=TRUE)

#Simulate Transition Probabilities When latency is high
set.seed(1009)
ms.hicap <- mssample(Haz = msf.hicap$Haz, trans = tmat, tvec = 0:40, history=list(state=2, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)


#Set nuclear latency to low
nuk7set1_1 <- as.vector(rep(0, 3))
nuk7set1_1[1] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .25)
nuk7set1_2 <- as.vector(rep(0, 3))
nuk7set1_2[2] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .25)
nuk7set1_3 <- as.vector(rep(0, 3))
nuk7set1_3[3] <- quantile(dat.1$nuk7set1, na.rm=TRUE, .25)

sit.locap <- data.frame(cbind(
							nuk7set1_1, lmid_past5yrs_1, npt_eff_1, neopatr_1,
							nuk7set1_2, lmid_past5yrs_2, npt_eff_2, neopatr_2,
							nuk7set1_3, lmid_past5yrs_3, npt_eff_3, neopatr_3))
							
sit.locap$trans <- 1:3							
attr(sit.locap, "trans") <- tmat
class(sit.locap) <- c("msdata", "data.frame")						
sit.locap$strata <- sit.locap$trans

#Estimate Cumulative Hazards When latency is low
msf.locap <- msfit(mod.1, sit.locap, trans=tmat, variance=TRUE)

#Simulate Transition Probabilities When latency is low
set.seed(1009)
ms.locap <- mssample(Haz = msf.locap$Haz, trans = tmat, tvec = 0:40, history=list(state=2, time=0), M = 10000, clock = "reset", output = "state", do.trace=1000)


#Figure 5a
plot(ms.locap$time, ms.locap$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,1), xlab="Years Since Program Initiated", ylab="Probability of Nuclear Program", main="")
lines(ms.locap$time, ms.locap$pstate2, lty=1)
lines(ms.hicap$time, ms.hicap$pstate2, lty=2)
legend(3,.4, c("Low Nuclear Latency","High Nuclear Latency"), lty = c(1,2), bty="n")

#Figure 5b
plot(ms.locap$time, ms.locap$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,.8), xlab="Years Since Program Initiated", ylab="Probability of No Nuclear Program or Weapon", main="")
lines(ms.locap$time, ms.locap$pstate1, lty=1)
lines(ms.hicap$time, ms.hicap$pstate1, lty=2)
legend(3,.7, c("Low Nuclear Latency","High Nuclear Latency"), lty = c(1,2), bty="n")

#Figure 5c
plot(ms.locap$time, ms.locap$pstate1, type="n" ,xlim=c(0,35), ylim=c(0,.04), xlab="Years Since Program Initiated", ylab="Probability of Nuclear Weapon", main="")
lines(ms.locap$time, ms.locap$pstate3, lty=1)
lines(ms.hicap$time, ms.hicap$pstate3, lty=2)
legend(-.8,.035, c("Low Nuclear Latency","High Nuclear Latency"), lty = c(1,2), bty="n")








